# -*- coding: utf-8 -*-

# @Time : 2022/3/20 14:31

# @Author : Aweo
# @File : qt_test.py
#import qdarkstyle
from PyQt5 import uic, QtWidgets
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtWidgets import QApplication
import cv2
from PyQt5.QtCore import Qt
from nano_predict import Read
from threading import Thread
import time
import os
import shutil

class Stats():
    def __init__(self):
        self.ui = uic.loadUi("ui_pro/my_ui1.0.ui")
        self.ui.BTop.clicked.connect(self.open_cam)
        self.ui.BTcl.clicked.connect(self.close_cam)

        self.ui.BT1.clicked.connect(self.start_detect)
        self.ui.BT1.setEnabled(False)
        self.ui.BT1.setStyleSheet("color: lightgray")

        self.ui.BT2.clicked.connect(self.close_detect)
        self.ui.BT2.setEnabled(False)
        self.ui.BT2.setStyleSheet("color: lightgray")

        self.ui.BT3.clicked.connect(self.save_singleIMG)
        self.ui.BT3.setEnabled(False)
        self.ui.BT3.setStyleSheet("color: lightgray")

        self.ui.BT4.clicked.connect(self.VIDEO_process)
        self.ui.BT4.setEnabled(False)
        self.ui.BT4.setStyleSheet("color: lightgray")

        self.ui.result.setText("识别未开启")
        bg = QPixmap('ui_pro/img/438868d5da02722066c147c4f072a695.png')
        self.ui.bg_img.setPixmap(bg)
        self.ui.frame.setText('请打开摄像头')
        self.ui.frame.setAlignment(Qt.AlignCenter)

        self.read = Read(label=['感染', '未感染'],
                         model_path='jetson_model/CNN5v2.pth',
                         model='CNN5')
        self.flag = False
        self.v1 = False



    def open_cam(self):
        self.capture = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        self.ui.BT1.setEnabled(True)
        self.ui.BT1.setStyleSheet("color: black")
        self.ui.BT2.setEnabled(True)
        self.ui.BT2.setStyleSheet("color: black")
        self.ui.BT3.setEnabled(True)
        self.ui.BT3.setStyleSheet("color: black")
        self.ui.BT4.setEnabled(True)
        self.ui.BT4.setStyleSheet("color: black")
        self.thread1 = Thread(target=self.cap_capture)
        self.thread1.start()


    def close_cam(self):
        self.f1 = False
        while self.thread1.is_alive():
            time.sleep(0.01)
        self.ui.BT1.setEnabled(True)
        self.ui.BT1.setStyleSheet("color: gray")
        self.ui.BT2.setEnabled(True)
        self.ui.BT2.setStyleSheet("color: gray")
        self.ui.BT3.setEnabled(True)
        self.ui.BT3.setStyleSheet("color: gray")
        self.ui.BT4.setEnabled(True)
        self.ui.BT4.setStyleSheet("color: gray")
        self.ui.frame.setText('请打开摄像头')
        self.ui.frame.setAlignment(Qt.AlignCenter)
        self.capture.release()

    def cap_capture(self):
        self.f1 = True
        while self.f1:
            _, self.img = self.capture.read(cv2.IMREAD_COLOR)
            self.imgshow = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
            self.imgshow = cv2.resize(self.imgshow, (420, 420))
            show_img = QImage(self.imgshow.data, self.imgshow.shape[0], self.imgshow.shape[1], QImage.Format_RGB888)
            self.ui.frame.setPixmap(QPixmap.fromImage(show_img))


    def cap_detect(self):
        self.flag = True
        while self.flag:
            res = self.read.Precdict(self.img)
            self.ui.result.setText(res)

    def close_detect(self):
        self.flag = False
        while self.thread2.is_alive():
            time.sleep(0.01)
        self.ui.result.setText("识别未开启")

    def start_detect(self):
        self.thread2 = Thread(target=self.cap_detect)
        self.thread2.start()


    def save_singleIMG(self):
        options = QtWidgets.QFileDialog.Options()
        options |= QtWidgets.QFileDialog.ReadOnly
        file_name, _ = QtWidgets.QFileDialog.getSaveFileName(None, "保存图像", "", "JPEG Files (*.jpg);;PNG Files (*.png)",
                                                             options=options)
        if file_name:
            cv2.imwrite(file_name, self.img)


    def VIDEO_process(self):
        self.v1 = not self.v1
        if self.v1:
            self.ui.BT4.setText("停止录制")
            self.thread3 = Thread(target=self.record)
            self.thread3.start()

        elif not self.v1:
           self.v1 = False
           self.ui.BT4.setText("开始录制")
           while self.thread3.is_alive():
               time.sleep(0.01)
           self.out.release()
           self.save_VIDEO()

    def record(self):
        frame_width = int(self.capture.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(self.capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        self.out = cv2.VideoWriter("tempFile.avi", fourcc, 30.0, (frame_width, frame_height))
        while self.v1:
            self.out.write(self.img)

    def save_VIDEO(self):
        options = QtWidgets.QFileDialog.Options()
        options |= QtWidgets.QFileDialog.ReadOnly
        video_name, _ = QtWidgets.QFileDialog.getSaveFileName(None, "保存视频", "", "AVI Files (*.avi);;MP4 Files (*.mp4)",
                                                             options=options)
        os.rename("tempFile.avi", video_name.split('/')[-1])
        shutil.move(video_name.split('/')[-1], video_name)



app = QApplication([])
stats = Stats()
# 显示ui
stats.ui.show()
app.exec_()
